from transformers import BertModel
import torch

#定义设备信息
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(DEVICE)

#加载预训练模型
pretrained = BertModel.from_pretrained(
    r"D:\AIModels\llm\bert-base-chinese\models--bert-base-chinese\snapshots\c30a6ed22ab4564dc1e3b2ecbf6e766b0611a33f"
    ).to(DEVICE)
#打印模型结构
# print(pretrained)

#定义下游任务（增量模型）
class Model(torch.nn.Module):
    """
    定义下游任务模型
    """
    def __init__(self):
        super().__init__()
        #设计全连接网络，实现二分类任务
        self.fc = torch.nn.Linear(768,2) #768是Bert模型的输出维度，2是二分类任务的类别数量

    def forward(self, input_ids, attention_mask, token_type_ids):
        """
        前向传播
        :param input_ids: 输入的文本id
        :param attention_mask: 输入的文本掩码
        :param token_type_ids: 输入的文本类型id
        :return: 输出的结果
        """
        #冻结Bert模型的参数，让其不参与训练
        with torch.no_grad(): #不计算梯度，节省内存和计算资源
            out = pretrained(
                input_ids=input_ids,           #输入的文本id
                attention_mask=attention_mask, #输入的文本掩码
                token_type_ids=token_type_ids  #输入的文本类型id
                )
        
        #增量模型参与训练
        out = self.fc(
            out.last_hidden_state[:,0]  #取[CLS]的输出作为增量模型的输入， [CLS]是Bert模型的特殊字符，代表整个句子的表示
            )
        return out
